"""
Cornel Movies Dialogs Corpus
https://www.cs.cornell.edu/~cristian/Cornell_Movie-Dialogs_Corpus.html
"""
import os
import logging

from . import utils

log = logging.getLogger("cornell")
DATA_DIR = "data/cornell"
# todo 语句之间的分割符号？
SEPARATOR = "+++$+++"


def load_dialogues(data_dir=DATA_DIR, genre_filter=''):
    """
    加载对话数据
    data_dir: 对话训练数据的路径
    genre_filte: 电影对话的电影类型
    :return: [电影对话列表]
    """
    movie_set = None # 电影id集
    if genre_filter:
        # 加载指定类型电影集
        movie_set = read_movie_set(data_dir, genre_filter)
        log.info("Loaded %d movies with genre %s", len(movie_set), genre_filter)
    log.info("Read and tokenise phrases...")
    # 加载电影短语
    lines = read_phrases(data_dir, movies=movie_set)
    log.info("Loaded %d phrases", len(lines))
    # 将电影短语组织成对话
    dialogues = load_conversations(data_dir, lines, movie_set)
    # 返回电影对白
    return dialogues


def iterate_entries(data_dir, file_name):
    """
    Iterate over entries in file
    按行读取数据集内容
    :param data_dir: 目录
    :param file_name: 数据集文件名
    :return:
    """
    with open(os.path.join(data_dir, file_name), "rb") as fd:
        for l in fd:
            # 按行读取文件数据
            l = str(l, encoding='utf-8', errors='ignore')
            # 先按照分隔符分割字符串，再去除分割后字符串左右两边的空白字符
            # map将上述两步骤的内容转换为list返回
            yield list(map(str.strip, l.split(SEPARATOR)))


def read_movie_set(data_dir, genre_filter):
    '''
    加载电影数据集
    return： 返回加载的电影id数据集
    '''
    res = set()
    # 加载movie_titles_metadata电影标题数据集
    for parts in iterate_entries(data_dir, "movie_titles_metadata.txt"):
        # 加载电影的id和对应的电影类型
        m_id, m_genres = parts[0], parts[5]
        # 进行电影类型过滤 todo 为什么要这步，是因为我们可以组织各种类型得的数据吗？
        if m_genres.find(genre_filter) != -1:
            # 只保存指定类型的电影
            res.add(m_id)
    return res


def read_phrases(data_dir, movies=None):
    '''
    加载电影对应的短语对话数据集
    return： 返回加载的电影短语数据集
    '''
    res = {} # 保存电影的对话短语，对话l_id为key，value为分割为单词后的句子单词list
    for parts in iterate_entries(data_dir, "movie_lines.txt"):
        # 对话的id、电影id，对话内容
        l_id, m_id, l_str = parts[0], parts[2], parts[4]
        # 防止加载了不是movies集中的短语对话信息
        if movies and m_id not in movies:
            continue
        # 分割句子为独立的单词
        tokens = utils.tokenize(l_str)
        if tokens:
            res[l_id] = tokens
    return res


def load_conversations(data_dir, lines, movies=None):
    '''
    return res:电影的对话列表，每个列表的元素对应多句对白
    '''
    res = []
    # 加载电影对话的数据集
    for parts in iterate_entries(data_dir, "movie_conversations.txt"):
        # 电影id，对话id列表（表示那些对话id是同一个场景下的对白）
        m_id, dial_s = parts[2], parts[3]
        # 仅加载电影集的电影对话信息
        if movies and m_id not in movies:
            continue
        # 去除[]字符，按照,分割句子
        l_ids = dial_s.strip("[]").split(", ")
        # 去除'字符，将结果组成list
        l_ids = list(map(lambda s: s.strip("'"), l_ids))
        # 将句子组织成对话dial，并存储在res中
        dial = [lines[l_id] for l_id in l_ids if l_id in lines]
        if dial:
            res.append(dial)
    return res


def read_genres(data_dir):
    res = {}
    for parts in iterate_entries(data_dir, "movie_titles_metadata.txt"):
        m_id, m_genres = parts[0], parts[5]
        l_genres = m_genres.strip("[]").split(", ")
        l_genres = list(map(lambda s: s.strip("'"), l_genres))
        res[m_id] = l_genres
    return res
